import backtrader as bt
import pandas as pd
import numpy as np

import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import talib as talib
import numpy as np


# 自定义板块数据加载
class SectorData(bt.feeds.PandasData):
    lines = ("funds_flow",)  # 资金流字段
    params = (
        ("datetime", None),
        ("funds_flow", 6),  # 假设第6列是资金流数据
    )


# 个股数据扩展字段
class StockData(bt.feeds.PandasData):
    lines = ("turnover", "market_cap")
    params = (
        ("turnover", 7),  # 换手率
        ("market_cap", 8),  # 流通市值
    )


class SectorSelector(bt.Indicator):
    """
    板块选择指标：计算过去5日资金流排名
    """

    lines = ("rank",)
    params = (("period", 5),)

    def __init__(self):
        self.ff_ma = bt.indicators.SMA(self.data.funds_flow, period=self.p.period)
        self.addminperiod(self.p.period)

    def next(self):
        # 获取全板块资金流数据（假设已加载所有板块数据）
        sector_flows = [d.ff_ma[0] for d in self.datas if isinstance(d, SectorData)]
        # 计算排名
        ranks = pd.Series(sector_flows).rank(ascending=False)
        self.l.rank[0] = ranks.get(self.data._name, 0)


class DragonStrategy(bt.Strategy):
    params = (
        ("sector_top", 3),  # 选前3板块
        ("stock_per_sector", 3),  # 每板块选3股
    )

    def __init__(self):
        # 板块选择指标
        self.sector_ranker = SectorSelector()
        # 个股评价指标
        self.stock_scorer = StockScorer()

    def next(self):
        # 1. 筛选热门板块
        sectors = [d for d in self.datas if isinstance(d, SectorData)]
        top_sectors = sorted(sectors, key=lambda x: self.sector_ranker[x].rank[0])[
            : self.p.sector_top
        ]

        # 2. 板块内选股
        selected_stocks = []
        for sector in top_sectors:
            stocks = self.get_sector_stocks(sector)
            scores = [(s, self.calc_stock_score(s)) for s in stocks]
            top_stocks = sorted(scores, key=lambda x: x[1], reverse=True)[
                : self.p.stock_per_sector
            ]
            selected_stocks.extend(top_stocks)

        # 3. 输出评级结果
        self.print_ratings(selected_stocks)

    def calc_stock_score(self, stock):
        """个股综合评分"""
        return (
            0.3 * self.momentum_score(stock)
            + 0.2 * self.liquidity_score(stock)
            + 0.5 * self.technical_score(stock)
        )

    def momentum_score(self, stock):
        """动量得分：20日涨幅 + 5日RSI"""
        ret = (stock.close[0] / stock.close[-20] - 1) * 100
        rsi = bt.indicators.RSI(stock.close, period=5)
        return 0.7 * ret + 0.3 * rsi[0]

    def liquidity_score(self, stock):
        """流动性得分：换手率 + 成交额"""
        turnover = stock.turnover[0]
        volume = stock.volume[0] * stock.close[0] / 1e8  # 亿元
        return 0.6 * turnover + 0.4 * volume

    def technical_score(self, stock):
        """技术面得分：均线 + MACD + 量价"""
        ma_score = 1 if stock.close[0] > stock.ma20[0] else 0
        macd_hist = bt.indicators.MACDHistogram(stock).histo[0]
        macd_score = 1 if macd_hist > 0 else -1
        volume_score = 1 if stock.volume[0] > stock.volume[-5:].mean() else 0
        return ma_score + macd_score + volume_score


class StockAnalyzer(bt.Observer):
    lines = ("ma5", "ma20", "macd", "rsi", "turnover")
    plotinfo = dict(plot=True, subplot=True)

    def __init__(self):
        # 计算所有技术指标
        self.l.ma5 = bt.indicators.SMA(self.data.close, period=5)
        self.l.ma20 = bt.indicators.SMA(self.data.close, period=20)
        self.l.macd = bt.indicators.MACDHistogram(self.data)
        self.l.rsi = bt.indicators.RSI(self.data.close)
        self.l.turnover = self.data.turnover


class RatingSystem(bt.Analyzer):
    def __init__(self):
        self.ratings = []

    def notify_order(self, order):
        pass  # 不需要处理订单

    def next(self):
        # 获取当天的评级数据
        for data in self.datas:
            if isinstance(data, StockData):
                score = self.strategy.calc_stock_score(data)
                rating = "A" if score > 80 else "B" if score > 60 else "C"
                self.ratings.append(
                    {
                        "date": self.datetime.date(),
                        "code": data._name,
                        "score": score,
                        "rating": rating,
                    }
                )

    def get_analysis(self):
        return pd.DataFrame(self.ratings)


if __name__ == "__main__":
    print("test")
    cerebro = bt.Cerebro()

    # 加载板块数据
    sectors = ["technology", "healthcare", "energy"]
    for sector in sectors:
        df = pd.read_csv(f"{sector}.csv")
        data = SectorData(dataname=df, name=sector)
        cerebro.adddata(data)

    # 加载个股数据（假设已按板块分类）
    stocks = ["tech_stock1", "tech_stock2", ..., "energy_stock3"]
    for stock in stocks:
        df = pd.read_csv(f"{stock}.csv")
        data = StockData(dataname=df, name=stock)
        cerebro.adddata(data)

    cerebro.addstrategy(DragonStrategy)
    cerebro.addobserver(StockAnalyzer)
    cerebro.addanalyzer(RatingSystem, _name="rating")
    results = cerebro.run()
    rating_report = results[0].analyzers.rating.get_analysis()
    print(rating_report.tail())
